package cn.langpy.recall.memory;

import cn.langpy.recall.common.ChatRecord;
import cn.langpy.recall.common.KnownRecord;
import cn.langpy.recall.common.MemorySpace;
import cn.langpy.recall.common.Response;
import cn.langpy.recall.core.Recall;
import io.weaviate.client.Config;
import io.weaviate.client.WeaviateAuthClient;
import io.weaviate.client.WeaviateClient;
import io.weaviate.client.base.Result;
import io.weaviate.client.base.util.Assert;
import io.weaviate.client.v1.auth.exception.AuthException;
import io.weaviate.client.v1.batch.api.ObjectsBatcher;
import io.weaviate.client.v1.batch.model.ObjectGetResponse;
import io.weaviate.client.v1.data.model.WeaviateObject;
import io.weaviate.client.v1.filters.Operator;
import io.weaviate.client.v1.filters.WhereFilter;
import io.weaviate.client.v1.graphql.model.GraphQLResponse;
import io.weaviate.client.v1.graphql.query.argument.HybridArgument;
import io.weaviate.client.v1.graphql.query.argument.NearTextArgument;
import io.weaviate.client.v1.graphql.query.argument.SortArgument;
import io.weaviate.client.v1.graphql.query.argument.SortArguments;
import io.weaviate.client.v1.graphql.query.builder.GetBuilder;
import io.weaviate.client.v1.graphql.query.fields.Field;
import io.weaviate.client.v1.graphql.query.fields.Fields;
import io.weaviate.client.v1.schema.model.Schema;
import io.weaviate.client.v1.schema.model.WeaviateClass;
import org.apache.commons.lang3.StringUtils;

import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.logging.Logger;

/**
 * weaviate实现方式
 *
 * @author zhangchang
 */
public class WeaviateMemory implements Recall {
    private static Logger log = Logger.getLogger(WeaviateMemory.class.getName());
    private WeaviateClient client;
    private String memoryUrl;
    private String memoryKey;
    private String embeddingUrl;
    private String embeddingModel;
    private String llmUrl;
    private String llmKey;
    private String llmModel;
    private final static String CHATHISTORY = "ChatHistory";
    private final static String KNOWNBASE = "Knownbase";
    private final static String vectorizer = "text2vec-ollama";

    public WeaviateClient getClient() {
        return client;
    }

    public void setClient(WeaviateClient client) {
        this.client = client;
    }

    public String getEmbeddingUrl() {
        return embeddingUrl;
    }

    public void setEmbeddingUrl(String embeddingUrl) {
        this.embeddingUrl = embeddingUrl;
    }

    public String getEmbeddingModel() {
        return embeddingModel;
    }

    public void setEmbeddingModel(String embeddingModel) {
        this.embeddingModel = embeddingModel;
    }

    public String getMemoryUrl() {
        return memoryUrl;
    }

    public void setMemoryUrl(String memoryUrl) {
        this.memoryUrl = memoryUrl;
    }

    public String getMemoryKey() {
        return memoryKey;
    }

    public void setMemoryKey(String memoryKey) {
        this.memoryKey = memoryKey;
    }

    public String getLlmUrl() {
        return llmUrl;
    }

    public void setLlmUrl(String llmUrl) {
        this.llmUrl = llmUrl;
    }

    public String getLlmKey() {
        return llmKey;
    }

    public void setLlmKey(String llmKey) {
        this.llmKey = llmKey;
    }

    public String getLlmModel() {
        return llmModel;
    }

    public void setLlmModel(String llmModel) {
        this.llmModel = llmModel;
    }

    @Override
    public boolean existsMemorySpace(String knowledgeName) {
        checkClient();
        Result<WeaviateClass> result1 = client.schema().classGetter().withClassName(knowledgeName).run();
        return result1.getResult() != null;
    }

    @Override
    public List<MemorySpace> listMemorySpaces() {
        checkClient();
        Result<Schema> result = client.schema().getter().run();
        List<WeaviateClass> classes = result.getResult().getClasses();
        List<MemorySpace> list = new ArrayList<>();
        classes.stream().forEach(c -> {
            MemorySpace space = new MemorySpace();
            space.setName(c.getClassName());
            space.setDescription(c.getDescription());
            Map<String, Map<String, Object>> moduleConfig = (Map<String, Map<String, Object>>)c.getModuleConfig();
            if (moduleConfig != null) {
                Map<String, Object> text2vecMap = moduleConfig.get(vectorizer);
                if (text2vecMap != null) {
                    space.setText2vecUrl(text2vecMap.get("apiEndpoint")+"");
                    space.setText2vecModel(text2vecMap.get("model")+"");
                }
            }

            list.add(space);
        });
        return list;
    }

    @Override
    public boolean deleteMemorySpace(String knowledgeName) {
        checkClient();
        Result<Boolean> result1 = client.schema().classDeleter().withClassName(knowledgeName).run();
        if (result1.getResult()) {
            log.info(String.format("【Memory】=> A memory space named %s was removed successfully.", knowledgeName));
        } else {
            log.info(String.format("【Memory】=> Can not remove memory space %s,beacause of '%s'.", knowledgeName, result1.getError().getMessages()));
        }
        return result1.getResult();
    }



    @Override
    public boolean createMemorySpace(String knowledgeName) {
        checkClient();
        checkText2vec();
        Map<String, Object> text2vecSettings = new HashMap<>();
        text2vecSettings.put("apiEndpoint", embeddingUrl);
        text2vecSettings.put("model", embeddingModel);

        Map<String, Map<String, Object>> moduleConfig = new HashMap<>();
        moduleConfig.put(vectorizer, text2vecSettings);

        WeaviateClass clazz = WeaviateClass.builder()
                .className(knowledgeName)
                .vectorizer(vectorizer)
                .moduleConfig(moduleConfig)
                .build();
        Result<Boolean> result = client.schema().classCreator().withClass(clazz).run();
        if (result.getResult()) {
            log.info(String.format("【Memory】=> A memory space named %s was created successfully.", knowledgeName));
        } else {
            log.info(String.format("【Memory】=> Can not create memory space %s,beacause of '%s'.", knowledgeName, result.getError().getMessages()));
        }
        return result.getResult();
    }


    private HashMap<String, Object> format(ChatRecord record) {
        HashMap<String, Object> properties = new HashMap<>();
        properties.put("role", record.getRole().getCode());
        properties.put("userId", record.getUserId());
        properties.put("sessionId", record.getSessionId());
        properties.put("message", record.getMessage());
        properties.put("createTime", LocalDateTime.now());
        return properties;
    }

    private HashMap<String, Object> format(KnownRecord record) {
        HashMap<String, Object> properties = new HashMap<>();
        properties.put("role", record.getRole().getCode());
        properties.put("userId", record.getUserId());
        properties.put("question", record.getQuestion());
        properties.put("answer", record.getAnswer());
        properties.put("createTime", LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss")));
        return properties;
    }

    private Result<ObjectGetResponse[]> add(String space, Map<String, Object> properties) {
        if (!existsMemorySpace(space)) {
            createMemorySpace(space);
        }
        ObjectsBatcher batcher = client.batch().objectsBatcher();
        batcher.withObject(WeaviateObject.builder()
                .className(space)
                .properties(properties)
                .build());
        Result<ObjectGetResponse[]> run = batcher.run();
        return run;
    }

    @Override
    public Response add(ChatRecord record) {
        Result<ObjectGetResponse[]> run = add(CHATHISTORY, format(record));
        Response response = new Response();
        if (run.hasErrors()) {
            response.setCode(0);
        }
        return response;
    }

    @Override
    public Response add(KnownRecord record) {
        String space = StringUtils.isBlank(record.getMemorySpace())?KNOWNBASE:record.getMemorySpace();
        Result<ObjectGetResponse[]> run = add(space, format(record));
        Response response = new Response();
        if (run.hasErrors()) {
            response.setCode(0);
        }
        return response;
    }


    @Override
    public Response call(ChatRecord record) {
        checkClient();
        List<WhereFilter> whereFilters = new ArrayList<>();
        if (!StringUtils.isBlank(record.getUserId())) {
            WhereFilter userWhere = WhereFilter.builder()
                    .path(new String[]{"userId"})
                    .operator(Operator.Equal)
                    .valueText(record.getUserId())
                    .build();
            whereFilters.add(userWhere);
        }
        if (!StringUtils.isBlank(record.getSessionId())) {
            WhereFilter userWhere = WhereFilter.builder()
                    .path(new String[]{"sessionId"})
                    .operator(Operator.Equal)
                    .valueText(record.getSessionId())
                    .build();
            whereFilters.add(userWhere);
        }
        Fields fields = Fields.builder()
                .fields(
                        Field.builder().name("userId").build(),
                        Field.builder().name("sessionId").build(),
                        Field.builder().name("message").build(),
                        Field.builder().name("createTime").build(),
                        Field.builder().name("_additional").fields(
                                Field.builder().name("distance").build(),
                                Field.builder().name("certainty").build()
                        ).build()
                ).build();

        WhereFilter where = WhereFilter.builder()
                .operator(Operator.And)
                .operands(whereFilters.toArray(new WhereFilter[whereFilters.size()]))
                .build();

        SortArguments sortArguments = SortArguments.builder().sort(new SortArgument[]{SortArgument.builder().path(new String[]{"createTime"}).build()}).build();
        GetBuilder.GetBuilderBuilder queryBuilder = GetBuilder.builder()
                .className(CHATHISTORY)
                .fields(fields)
                .withWhereFilter(where)
                .withSortArguments(sortArguments)
                .limit(100);

        if (!StringUtils.isBlank(record.getMessage())) {
            NearTextArgument nearText = NearTextArgument.builder()
                    .concepts(new String[]{record.getMessage()})
                    .certainty(0.8f)
                    .build();
            queryBuilder.withNearTextFilter(nearText);
        }

        Result<GraphQLResponse> result = client.graphQL().raw().withQuery(queryBuilder.build().buildQuery()).run();
        Response response = new Response();
        if (result.hasErrors()) {
            response.setCode(0);
        } else {
            response.setData(result.getResult().getData());
        }
        return response;
    }


    @Override
    public Response call(KnownRecord record) {
        checkClient();
        List<String> queries = new ArrayList<>();
        if (!StringUtils.isBlank(record.getUserId())) {
            queries.add("userId");
        }
        if (!StringUtils.isBlank(record.getQuestion())) {
            queries.add("question");
        }
        HybridArgument hybridArgument = HybridArgument.builder()
                .query(record.getQuestion())
                .properties(queries.toArray(new String[queries.size()]))
                .alpha(0.7f)
                .build();

        Fields fields = Fields.builder()
                .fields(
                        Field.builder().name("userId").build(),
                        Field.builder().name("question").build(),
                        Field.builder().name("answer").build(),
                        Field.builder().name("createTime").build(),
                        /*返回相似度*/
                        Field.builder().name("_additional").fields(
                                /*确定性*/
                                Field.builder().name("certainty").build()
                        ).build()
                ).build();

        String query = GetBuilder.builder()
                .className(KNOWNBASE)
                .fields(fields)
                .withHybridFilter(hybridArgument)
                .limit(5)
                .build()
                .buildQuery();

        Result<GraphQLResponse> result = client.graphQL().raw().withQuery(query).run();
        Response response = new Response();
        if (result.hasErrors()) {
            response.setCode(0);
        } else {
            response.setData(result.getResult().getData());
        }
        return response;
    }

    public static WeaviateMemoryBuilder builder() {
        return new WeaviateMemoryBuilder();
    }

    private void checkClient() {
        Assert.requiredNotNull(this.client, "WeaviateClient");
    }

    private void checkText2vec() {
        Assert.requiredNotNull(this.embeddingUrl, "text2vecUrl");
        Assert.requiredNotNull(this.embeddingModel, "text2vecModel");
    }

    public static class WeaviateMemoryBuilder {
        private WeaviateMemory memory;

        public WeaviateMemoryBuilder() {
            memory = new WeaviateMemory();
        }

        public WeaviateMemory build() {
            if (memory.getClient() == null) {
                if (!StringUtils.isBlank(memory.getMemoryUrl())){
                    String[] split = memory.getMemoryUrl().split("://");
                    if (split.length != 2) {
                        throw new RuntimeException("WeaviateMemory requires a valid memory url like 'http://localhost:8080'");
                    }
                    Config config = new Config(split[0], split[1]);
                    WeaviateClient client = null;
                    if (!StringUtils.isBlank(memory.getMemoryKey())) {
                        try {
                            client = WeaviateAuthClient.apiKey(config,memory.getMemoryKey());
                        } catch (AuthException e) {
                            throw new RuntimeException(e);
                        }
                    }else {
                        client = new WeaviateClient(config);
                    }
                    memory.setClient(client);
                }
            }
            return memory;
        }

        public WeaviateMemoryBuilder client(WeaviateClient client) {
            this.memory.setClient(client);
            return this;
        }

        public WeaviateMemoryBuilder embeddingUrl(String text2vecUrl) {
            this.memory.setEmbeddingUrl(text2vecUrl);
            return this;
        }

        public WeaviateMemoryBuilder embeddingModel(String text2vecModel) {
            this.memory.setEmbeddingModel(text2vecModel);
            return this;
        }
        public WeaviateMemoryBuilder memoryUrl(String memoryUrl) {
            this.memory.setMemoryUrl(memoryUrl);
            return this;
        }
        public WeaviateMemoryBuilder memoryKey(String memoryKey) {
            this.memory.setMemoryKey(memoryKey);
            return this;
        }
        public WeaviateMemoryBuilder llmUrl(String llmUrl) {
            this.memory.setLlmUrl(llmUrl);
            return this;
        }
        public WeaviateMemoryBuilder llmModel(String llmModel) {
            this.memory.setLlmModel(llmModel);
            return this;
        }
        public WeaviateMemoryBuilder llmKey(String llmKey) {
            this.memory.setLlmKey(llmKey);
            return this;
        }
    }

}

